#!/usr/bin/env python

from pywolb.deterministic import base, meta, weights
from pywolb.utilities import numpytools, stringtools, diversetools
import numpy as np
import copy, os, gzip, cPickle, pdb

for mod in [meta,base,weights,numpytools,stringtools]:
    reload(mod)

def line(char='-', n=70):
    return n*char

def main():
    global metapop
        
    try:
        
#        ip1 = 0.01          # introduction frequency of mutant allele P1
        m1  = 0.05           # migration rate
        m2  = 0.1
        pr  = 1.             # rejection probability
        pt  = 0.8            # transition probability
        r   = 0.5            # recombination rate between trait and preference locus
        s   = 0.25           # selection coefficients for populations (1, 2)
        s1  = [s, 0., 0., 0] # A1 and B1 favored separately in left population
        s2  = [0., s, s, 0.] # A2,A3 and B2,B3 favored in middle population/hybrid zone
        s3  = [0., 0., 0.,s] # A4 and B4 favored in right population
                             # selection coefficients at the A and B loci are multiplicative ,
                             # i.e. (1+s)*(1+s) !
        ci  = 0.9
        e   = 1.             # penetrance of the CI suppressor
        u   = 0.05           # viability cost to males carrying the suppressor
        f   = 0.
        t   = 0.9
        h   = 1.             # sterility coefficient (HMS)
        d   = 1.             # factor for partial HMS
        
        print line()
        print "Parameters: m1=%s, m2=%s, ci=%s, e=%s, u=%s, f=%s, h=%s, d=%s, t=%s, pr=%s, pt=%s, r=%s, s=%s" % (str(m1), str(m2), str(ci), str(e), str(u), str(f), str(h), str(d), str(t), str(pr), str(pt), str(r), str(s))

        freqs1 = np.zeros((3,4,4,2,3,2))
        freqs1[0,0,0,0,0,0] = 0.5    # Y0-A1-B1-P0-S0-Uninf (females) 
        freqs1[1,0,0,0,0,0] = 0.5    # Y1-A1-B1-P0-S0-Uninf (males)
        
        freqs2 = np.zeros((3,4,4,2,3,2))
        freqs2[0,1,1,1,0,0] = 0.25   # Y0-A2-B2-P1-S0-Uninf (females)   
        freqs2[1,1,1,1,0,0] = 0.25   # Y1-A2-B2-P1-S0-Uninf (males)
        freqs2[0,2,2,0,0,1] = 0.25   # Y0-A3-B3-P0-S0-Wolb (females)   
        freqs2[2,2,2,0,0,1] = 0.25   # Y2-A2-B2-P1-S0-Uninf (males)
        
        freqs3 = np.zeros((3,4,4,2,3,2))
        freqs3[0,3,3,0,0,1] = 0.5    # Y0-A4-B4-P0-S0-Wolb (females)
        freqs3[2,3,3,0,0,1] = 0.5    # Y2-A4-B4-P0-S0-Wolb (males)
        
       
        labels = {'axes': ['y-chromosome', 'background-A', 'background-B', 'preference', 'sterility-marker', 'cytotype'], \
                  'elements': [['Y0','Y1','Y2'], \
                               ['A1','A2','A3','A4'], \
                               ['B1','B2','B3','B4'], \
                               ['P0','P1'], \
                               ['S0','S1','S2'], \
                               ['Uninf', 'Wolb']]}
        
        CI   = weights.YModifiedCI(ci_level=ci, mod_penetrance=e)
        F    = weights.FecundityReduction(fecundity_reduction=f)
        T    = weights.WolbachiaTransmission(transmission_rate=t)
#        HMS  = weights.HybridMaleSterility2CustomLoci4Alleles(sterility_coefficient=h, locus1=('background-A','A'), locus2=('background-B','B'))
        MP   = weights.MatingPreferenceAB(pt, 'pr', 'P0','None',0., 'P1','A2-B2',pr)
        NIP  = weights.NuclearInheritanceAtSingleLocus('preference', 'P', range(0,2))
        NIAB = weights.NuclearInheritanceAtTwoLoci(r, 'background-A', 'A', range(1,5), 'background-B', 'B', range(1,5))
        NIY  = weights.NuclearInheritanceAtYLocus(2)
        RC   = weights.ResistanceCostsY(viability_reduction=u)
        MF  =  weights.MaleFertility(sterility_coefficient=h, partial_factor=d)
        HMS  = weights.SetOffspringFertility()
        
        pop1 = base.Population(freqs1, labels)
        pop1.add_weights(CI, RC, F, T, MP, NIY, NIAB, NIP, MF, HMS)
        pop1.add_weights(weights.ViabilitySelectionAtCustomLocus('background-A', s1), \
                         weights.ViabilitySelectionAtCustomLocus('background-B', s1))   
        
        pop2 = base.Population(freqs2, labels)
        pop2.add_weights(CI, RC, F, T, MP, NIY, NIAB, NIP, MF, HMS)
        pop2.add_weights(weights.ViabilitySelectionAtCustomLocus('background-A', s2), \
                         weights.ViabilitySelectionAtCustomLocus('background-B', s2))
        
        pop3 = base.Population(freqs3, labels)
        pop3.add_weights(CI, RC, F, T, MP, NIY, NIAB, NIP, MF, HMS)
        pop3.add_weights(weights.ViabilitySelectionAtCustomLocus('background-A', s3), \
                         weights.ViabilitySelectionAtCustomLocus('background-B', s3))
        

        # stepping stones:
#        migrates = np.array([[1-m,     m,   0], \
#                             [m,   1-2*m,   m], \
#                             [0.,      m, 1-m]])
        # two mainlands and hybrid zone:
#        migrates = np.array([[1,       0,   0], \
#                             [m,   1-2*m,   m], \
#                             [0,       0,   1]])
        # asymmetrical migration:
        migrates = np.array([[1-m1,       m1,      0], \
                             [m1,   1-m1-m2,      m2], \
                             [0,          m2,   1-m2]])
        metapop = meta.MetaPopulation(migrates, pop1, pop2, pop3)
        print line()
        print "Starting condition"
        print line()
        print metapop._show('full overview', 10)

        prec = 1e-4
#        metapop.run(1000)
        metapop.run(1)
        print line()
        print metapop._show('full overview', 10)
            
#        metapop.reset()
#        print "Introducing mating preference"
#        metapop.introduce_alleles('pop2','BP1',ip1, 'pop3','BP2',ip1)
#        metapop.show_overview()
#        
#        eq_found = metapop.findEQ(prec,nmax=5e4)
#        print "Final"
#        metapop.show('overview')

#        
#        metapop.calculate_effective_migration_matrix(n=100)
#        metapop.calculate_gene_flow_factors()
#        mAeff,mBeff = metapop.effM[0,1],metapop.effM[1,0]
#        gffA,gffB = metapop.GFF[0,1],metapop.GFF[1,0]
#        # most frequent types (label,frequency):
#        lA,xA = metapop.pop1.get_most_frequent()[0]
#        lB,xB = metapop.pop2.get_most_frequent()[0]
#        ls = len(lA)        # length of type string
#        popwidth = max(ls+3,9)
#        middlewidth = 18
#        print line()
#        print "Gene flow factors (effective migration)"
#        print line()
#        print "1%s %s %s2" % ((popwidth-1)*'_', middlewidth*' ', (popwidth-1)*'_')
#        print "%s|%s|%s"   % (popwidth*' ', middlewidth*' ', popwidth*' ')
#        print "%s|  <--[%5.1f%%]--   |%s" % (lA.center(popwidth), 100*gffA, lB.center(popwidth))
#        print "%s|   --[%5.1f%%]-->  |%s" % ("%5.1f%%".center(popwidth+1)%(100*xA), 100*gffB, "%5.1f%%".center(popwidth+1)%(100*xB))
#        print "%s|%s|%s"   % (popwidth*'_', middlewidth*' ', popwidth*'_')
    
#    except:
#        raise
    finally:
        return metapop


if __name__ == '__main__':
    mp = main()

